package top.wshape1.rpc.loadbalance;

import top.wshape1.rpc.common.entity.ServiceMeta;
import top.wshape1.rpc.config.JJRpcConfig;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;

/**
 * refer to guide-rpc-framework consistent hash load balance:
 *  <a href="https://gitee.com/SnailClimb/guide-rpc-framework/blob/master/rpc-framework-simple/src/main/java/github/javaguide/loadbalance/loadbalancer/ConsistentHashLoadBalance.java">Guide_ConsistentHashLoadBalance.java</a>
 *  and dubbo_v2.7.x consistent hash load balance:
 *  <a href="https://github.com/apache/dubbo/blob/2.7.x/dubbo-cluster/src/main/java/org/apache/dubbo/rpc/cluster/loadbalance/ConsistentHashLoadBalance.java">Dubbo_ConsistentHashLoadBalance.java</a>
 *
 *
 * @author Wshape1
 * @version 1.0
 * @description ConsistentHashLoadBalance 一致性哈希
 * @since 2024-02-10
 */

public class ConsistentHashLoadBalance extends AbstractLoadBalance {

    private static final int DEFAULT_REPLICA_NUMBER = 160;

    private final ConcurrentHashMap<String, ConsistentHashSelector> selectors = new ConcurrentHashMap<>();

    @Override
    protected ServiceMeta doSelect(List<ServiceMeta> serviceMetas) {
        // 获取服务列表的HashCode，判断服务列表是否更新
        int identityHashCode = Objects.hashCode(serviceMetas);

        // 获取调用的服务名
        String serviceName = serviceMetas.get(0).getServiceName();

        // 获取一致性哈希选择器
        ConsistentHashSelector selector = selectors.get(serviceName);
        // 判断一致性哈希选择器是否存在
        if (selector == null || selector.identityHashCode != identityHashCode) {
            selectors.put(serviceName, new ConsistentHashSelector(serviceMetas, DEFAULT_REPLICA_NUMBER, identityHashCode));
            selector = selectors.get(serviceName);
        }
        // 获取调用的服务地址
        String myAddress = JJRpcConfig.CONFIG.getHost() + ":" + JJRpcConfig.CONFIG.getPort();
        return selector.select(myAddress);
    }

    private static class ConsistentHashSelector {
        private final TreeMap<Long, ServiceMeta> virtualInvokers;

        private final int identityHashCode;

        ConsistentHashSelector(List<ServiceMeta> invokers, int replicaNumber, int identityHashCode) {
            this.virtualInvokers = new TreeMap<>();
            this.identityHashCode = identityHashCode;

            for (ServiceMeta invoker : invokers) {
                String address = invoker.getAddress();
                for (int i = 0; i < replicaNumber / 4; i++) {
                    byte[] digest = md5(address + i);
                    for (int h = 0; h < 4; h++) {
                        long m = hash(digest, h);
                        virtualInvokers.put(m, invoker);
                    }
                }
            }
        }

        private static byte[] md5(String key) {
            MessageDigest md;
            try {
                md = MessageDigest.getInstance("MD5");
                byte[] bytes = key.getBytes(StandardCharsets.UTF_8);
                md.update(bytes);
            } catch (NoSuchAlgorithmException e) {
                throw new IllegalStateException(e.getMessage(), e);
            }

            return md.digest();
        }

        private static long hash(byte[] digest, int idx) {
            return ((long) (digest[3 + idx * 4] & 255) << 24 | (long) (digest[2 + idx * 4] & 255) << 16 | (long) (digest[1 + idx * 4] & 255) << 8 | (long) (digest[idx * 4] & 255)) & 4294967295L;
        }

        public ServiceMeta select(String address) {
            byte[] digest = md5(address);
            return selectForKey(hash(digest, 0));
        }

        public ServiceMeta selectForKey(long hashCode) {
            Map.Entry<Long, ServiceMeta> entry = virtualInvokers.ceilingEntry(hashCode);

            if (entry == null) {
                entry = virtualInvokers.firstEntry();
            }

            return entry.getValue();
        }
    }
}
